from concurrent.futures import ThreadPoolExecutor, ProcessPoolExecutor
from typing import Callable, Iterable, Any
import logging


class ParallelExecutor:
    def __init__(self,
                 max_workers: int = None,
                 pool_type: str = 'thread',
                 auto_shutdown: bool = True):
        self.pool_type = pool_type.lower()
        self.executor = (ThreadPoolExecutor if pool_type == 'thread'
                         else ProcessPoolExecutor)(max_workers=max_workers)
        self.auto_shutdown = auto_shutdown

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        if self.auto_shutdown:
            self.executor.shutdown(wait=True)

    def batch_submit(self,
                     tasks: list[tuple[Callable, tuple]],
                     callbacks: list[Callable[[Any], None]] = None) -> list:
        futures = []
        try:
            for index, (func, args) in enumerate(tasks):
                future = self.executor.submit(func, *args)
                if callbacks and index < len(callbacks):
                    future.add_done_callback(
                        lambda f, cb=callbacks[index]: cb(f.result())
                    )
                futures.append(future)
            return futures
        except Exception as e:
            logging.error(f"Multi-task submit failed: {str(e)}")
            raise

    @staticmethod
    def wait_all(futures: list, timeout: float = None):
        results = []
        for future in futures:
            try:
                results.append(future.result(timeout))
            except Exception as e:
                results.append(e)
        return results


# 使用示例
if __name__ == "__main__":
    # 定义不同任务
    def task1(x):
        return x * 2


    def task2(y):
        return y + 10


    # 回调函数
    def callback1(result):
        print(f"Task1 result: {result}")


    def callback2(result):
        print(f"Task2 result: {result}")


    # 提交混合任务
    with ParallelExecutor() as executor:
        futures = executor.batch_submit(
            tasks=[
                (task1, (5,)),
                (task2, (3,)),
                (lambda z: z ** 3, (4,))
            ],
            callbacks=[callback1, callback2]
        )
        print("All results:", executor.wait_all(futures))
